import pickle
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from collections import Counter
import os


def load_and_analyze_graph(pickle_file):
    """Load and analyze graph structure"""
    try:
        with open(pickle_file, "rb") as f:
            G = pickle.load(f)

        print(f"Successfully loaded graph file: {pickle_file}")
        print(f"Graph type: {type(G)}")

        if isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph)):
            print(f"Number of nodes: {G.number_of_nodes()}")
            print(f"Number of edges: {G.number_of_edges()}")
            print(f"Is directed: {G.is_directed()}")
            print(f"Is multigraph: {G.is_multigraph()}")

            # Analyze node attributes
            if G.nodes():
                sample_node = list(G.nodes())[0]
                print(f"Sample node: {sample_node}")
                print(f"Node attributes: {G.nodes[sample_node]}")

            # Analyze edge attributes - Fixed MultiDiGraph edge access
            if G.edges():
                # For MultiGraph/MultiDiGraph, need to explicitly get edges with keys
                if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
                    # Get edges with keys
                    edges_with_keys = list(G.edges(keys=True))
                    if edges_with_keys:
                        sample_edge = edges_with_keys[0]
                        print(f"Sample edge: {sample_edge}")
                        u, v, key = sample_edge
                        edge_attrs = G.edges[u, v, key]
                        if edge_attrs:
                            print(f"Edge attributes: {edge_attrs}")
                        else:
                            print("Edge attributes: {}")
                else:
                    # Regular graph handling
                    sample_edge = list(G.edges())[0]
                    print(f"Sample edge: {sample_edge}")
                    edge_attrs = G.edges[sample_edge]
                    if edge_attrs:
                        print(f"Edge attributes: {edge_attrs}")
                    else:
                        print("Edge attributes: {}")

            return G
        else:
            print(f"Not a NetworkX graph object, but: {type(G)}")
            return G

    except Exception as e:
        print(f"Error loading file: {e}")
        return None


def visualize_graph_basic(G, output_dir="demo/visualizations"):
    """Basic graph visualization"""
    if not isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph)):
        print("Not a NetworkX graph object, cannot visualize")
        return

    os.makedirs(output_dir, exist_ok=True)

    plt.figure(figsize=(15, 10))

    # If too many nodes, only show a subset
    if G.number_of_nodes() > 100:
        # Select nodes with highest degree
        degrees = dict(G.degree())
        top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:50]
        subgraph = G.subgraph([node for node, _ in top_nodes])
        print(
            f"Too many nodes ({G.number_of_nodes()}), showing only top 50 nodes by degree"
        )
    else:
        subgraph = G

    try:
        # Use spring layout
        pos = nx.spring_layout(subgraph, k=1, iterations=50)

        # Set different colors based on node type
        node_colors = []
        for node in subgraph.nodes():
            node_type = subgraph.nodes[node].get("type", "unknown")
            if node_type == "directory":
                node_colors.append("lightcoral")
            elif node_type == "file":
                node_colors.append("lightblue")
            elif node_type == "function":
                node_colors.append("lightgreen")
            elif node_type == "class":
                node_colors.append("lightyellow")
            else:
                node_colors.append("lightgray")

        # Draw nodes
        nx.draw_networkx_nodes(
            subgraph, pos, node_size=50, node_color=node_colors, alpha=0.7
        )

        # Draw edges
        nx.draw_networkx_edges(subgraph, pos, alpha=0.3, width=0.5, edge_color="gray")

        # If not too many nodes, show labels
        if subgraph.number_of_nodes() <= 20:
            # Truncate long labels
            labels = {}
            for node in subgraph.nodes():
                label = str(node)
                if len(label) > 15:
                    label = label[:12] + "..."
                labels[node] = label
            nx.draw_networkx_labels(subgraph, pos, labels, font_size=6)

        plt.title(
            f"Dependency Graph Visualization (Nodes: {subgraph.number_of_nodes()}, Edges: {subgraph.number_of_edges()})"
        )
        plt.axis("off")
        plt.tight_layout()
        plt.savefig(f"{output_dir}/graph_basic.png", dpi=300, bbox_inches="tight")
        plt.show()

    except Exception as e:
        print(f"Error during visualization: {e}")


def analyze_graph_statistics(G, output_dir="visualizations"):
    """Analyze graph statistics"""
    if not isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph)):
        return

    os.makedirs(output_dir, exist_ok=True)

    # Degree distribution
    degrees = [d for n, d in G.degree()]

    plt.figure(figsize=(15, 5))

    # Degree distribution histogram
    plt.subplot(1, 3, 1)
    plt.hist(degrees, bins=min(50, len(set(degrees))), alpha=0.7)
    plt.xlabel("Degree")
    plt.ylabel("Number of Nodes")
    plt.title("Degree Distribution")
    plt.yscale("log")

    # Cumulative degree distribution
    plt.subplot(1, 3, 2)
    degree_counts = Counter(degrees)
    degrees_sorted = sorted(degree_counts.keys())
    cumulative = np.cumsum([degree_counts[d] for d in degrees_sorted])
    plt.plot(degrees_sorted, cumulative)
    plt.xlabel("Degree")
    plt.ylabel("Cumulative Number of Nodes")
    plt.title("Cumulative Degree Distribution")
    plt.yscale("log")

    # Connected component sizes
    plt.subplot(1, 3, 3)
    if G.is_directed():
        components = list(nx.weakly_connected_components(G))
    else:
        components = list(nx.connected_components(G))

    component_sizes = [len(c) for c in components]
    plt.hist(component_sizes, bins=min(20, len(set(component_sizes))), alpha=0.7)
    plt.xlabel("Component Size")
    plt.ylabel("Number of Components")
    plt.title("Connected Component Size Distribution")

    plt.tight_layout()
    plt.savefig(f"{output_dir}/graph_statistics.png", dpi=300, bbox_inches="tight")
    plt.show()

    # Print statistics
    print(f"\n=== Graph Statistics ===")
    print(f"Average degree: {np.mean(degrees):.2f}")
    print(f"Maximum degree: {max(degrees)}")
    print(f"Minimum degree: {min(degrees)}")
    print(f"Number of connected components: {len(components)}")
    print(f"Largest component size: {max(component_sizes) if component_sizes else 0}")


def analyze_node_types(G, output_dir="visualizations"):
    """Analyze node types"""
    if (
        not isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph))
        or not G.nodes()
    ):
        return

    # Analyze node types (based on node attributes)
    node_types = {}
    for node in G.nodes():
        node_type = G.nodes[node].get("type", "unknown")
        node_types[node] = node_type

    type_counts = Counter(node_types.values())
    print(f"\n=== Node Type Analysis ===")
    for node_type, count in type_counts.items():
        print(f"{node_type}: {count}")

    # Visualize node type distribution
    if len(type_counts) > 1:
        plt.figure(figsize=(10, 6))
        plt.pie(type_counts.values(), labels=type_counts.keys(), autopct="%1.1f%%")
        plt.title("Node Type Distribution")
        plt.savefig(f"{output_dir}/node_types.png", dpi=300, bbox_inches="tight")
        plt.show()

    # Analyze file extensions
    file_extensions = {}
    for node in G.nodes():
        if G.nodes[node].get("type") == "file":
            node_str = str(node)
            if "." in node_str:
                ext = node_str.split(".")[-1]
                file_extensions[ext] = file_extensions.get(ext, 0) + 1

    if file_extensions:
        print(f"\n=== File Extension Analysis ===")
        for ext, count in sorted(
            file_extensions.items(), key=lambda x: x[1], reverse=True
        ):
            print(f".{ext}: {count}")


def analyze_directory_structure(G):
    """Analyze directory structure"""
    print(f"\n=== Directory Structure Analysis ===")

    # Find all directory nodes
    directories = [
        node for node in G.nodes() if G.nodes[node].get("type") == "directory"
    ]
    files = [node for node in G.nodes() if G.nodes[node].get("type") == "file"]

    print(f"Number of directories: {len(directories)}")
    print(f"Number of files: {len(files)}")

    # Analyze directory depth
    directory_depths = {}
    for directory in directories:
        depth = str(directory).count("/") if str(directory) != "/" else 0
        directory_depths[directory] = depth

    if directory_depths:
        max_depth = max(directory_depths.values())
        print(f"Maximum directory depth: {max_depth}")

        # Count directories by depth
        depth_counts = Counter(directory_depths.values())
        print("Directory count by depth:")
        for depth in sorted(depth_counts.keys()):
            print(f"  Depth {depth}: {depth_counts[depth]} directories")


def main():
    pickle_file = "/data_ext/ref_code/LocAgent/demo/data/astropy_astropy.pkl"

    # Check if file exists
    if not os.path.exists(pickle_file):
        print(f"File does not exist: {pickle_file}")
        return

    print(f"Analyzing file: {pickle_file}")
    print("=" * 50)

    # Load and analyze graph
    G = load_and_analyze_graph(pickle_file)

    if G is None:
        return

    if isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph)):
        # Create visualizations
        print("\nStarting visualization...")
        visualize_graph_basic(G)
        analyze_graph_statistics(G)
        analyze_node_types(G)
        analyze_directory_structure(G)

        print(f"\nVisualization complete! Images saved in visualizations/ directory")
    else:
        print(f"Object type: {type(G)}")
        if hasattr(G, "__len__"):
            print(f"Object length: {len(G)}")
        if hasattr(G, "keys"):
            print(f"Keys: {list(G.keys())[:10]}...")  # Only show first 10 keys


if __name__ == "__main__":
    main()
